import numpy as np
import scipy
import time
import ConvertPLY
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import scipy.ndimage as ndi
import scipy.io
import sys

import numpy.linalg as la
import combiner2 as combiner

def getelipse(N=100., phi = 0, rot = .3, t = [1.,.5, .25]):
  theta = np.pi*np.arange(N)*1./N
  B = np.zeros((3,N))

  B[0,:] = 2*np.cos(theta)
  B[1,:] = np.sin(theta)
  B[2,:] = 0

  A = np.array(B)
  
  A[0,:] = 2*np.cos(theta-phi)
  A[1,:] = np.sin(theta-phi)
  A[2,:] = 0
  Aori = np.array(A)
  rotmat = np.array([[np.cos(rot), -np.sin(rot), 0],
                        [np.sin(rot), np.cos(rot),0],
                         [0 , 0, 1.]] )
  
  A = np.dot(rotmat,A) + np.array(t)[:,None]
  return A,B, Aori

def getcubiceasy(N=300., phi = 0, rot = .3, t = [-1.,.5, .25]):
  x = np.arange(N)/(N/2) - .8
  B = np.zeros((3,N))

  B[0,:] = x
  B[1,:] = x**3
  B[2,:] = 0
  
  x = np.arange(N+100)/(N/2) - .9
  A = np.zeros((3,N+100))

  A[0,:] = x
  A[1,:] = x**3
  A[2,:] = 0
  Aorig = np.array(A)
  A += np.random.randn(3,N+100)*0.01
  B += np.random.randn(3,N)*0.01
  
  rot = np.pi/4
  rotmat = np.array([[np.cos(rot), -np.sin(rot), 0],
                        [np.sin(rot), np.cos(rot),0],
                         [0 , 0, 1.]] )
  
  A = np.dot(rotmat,A) + np.array(t)[:,None]
  randind = (np.random.rand(N/3)*N).astype(int)
  #A = A[:,randind]
  return A,B,Aorig


# harder test case, neither A includes B nor B includes A
def getcubichard(N=300., phi = 0, rot = .3, t = [-1.,.5, .25]):
  x = np.arange(N)/(N/2) - .95
  B = np.zeros((3,N))

  B[0,:] = x
  B[1,:] = x**3
  B[2,:] = 0
  
  x = np.arange(N)/(N/2) - 1.05
  A = np.zeros((3,N))

  A[0,:] = x
  A[1,:] = x**3
  A[2,:] = 0
  Aorig = np.array(A)
  A += np.random.randn(3,N)*0.01
  B += np.random.randn(3,N)*0.01
  
  rot = np.pi/3
  rotmat = np.array([[np.cos(rot), -np.sin(rot), 0],
                        [np.sin(rot), np.cos(rot),0],
                         [0 , 0, 1.]] )
  
  A = np.dot(rotmat,A) + np.array(t)[:,None]
  randind = (np.random.rand(N/3)*N).astype(int)
  #A = A[:,randind]
  return A,B,Aorig


# harder test case, L shaped thing
def getL(N=300., phi = 0, rot = .3, t = [-1.,.5, .25]):
  x = np.arange(3*N)*1./(N)
  x[2*N:3*N] = 2.
  y = np.zeros(3*N)
  y[2*N:3*N] = np.arange(N)*1./N

  B = np.zeros((3,3*N))

  B[0,:] = x
  B[1,:] = y
  B[2,:] = 0
  
  x = np.arange(2*N)*1./N
  A = np.zeros((3,2*N))

  A[0,:] = x
  A[1,:] = 0
  A[2,:] = 0
  Aorig = np.array(A)
  A += np.random.randn(*A.shape)*0.01
  B += np.random.randn(*B.shape)*0.01
  
  rot = np.pi/3
  rotmat = np.array([[np.cos(rot), -np.sin(rot), 0],
                        [np.sin(rot), np.cos(rot),0],
                         [0 , 0, 1.]] )
  
  A = np.dot(rotmat,A) + np.array(t)[:,None]
  randind = (np.random.rand(N/3)*N).astype(int)
  #A = A[:,randind]
  return A,B,Aorig

datafunctions = [getelipse, getcubiceasy, getcubichard, getL]

def testfakedata(ind = 0):
  A,B,Ao = datafunctions[ind].__call__()
  R,t = combiner.probcombine(A,B)
  print R
  print t
  C = np.dot(R,B) + t
  import matplotlib.pyplot as plt
  plt.scatter(A[0,:],A[1,:], c='r', alpha = .5)
  plt.scatter(B[0,:],B[1,:], c='g', alpha = .5)
  plt.scatter(C[0,:],C[1,:], c='b', alpha = .5)
  plt.scatter(Ao[0,:],Ao[1,:], c='b', marker = '+',alpha = .5)
  plt.axis([-3, 3, -3, 3])

firstplot = True
def plotfakedata(ind = 2, maxiter = 500, interval = 50):
  global Rs, ts, firstplot 
  A,B,Ao = datafunctions[ind].__call__()
  if firstplot:
    R,t = combiner.probcombine(A,B)
    Rs = combiner.Rs
    ts = combiner.ts
    firstplot = False
  select = np.arange(0,maxiter, interval)
  
  for s in select:
    print s
    R = Rs[s]; t = ts[s][:,None]
    if la.det(R) < 0.1: break
    C = np.dot(R,B) + t
    import matplotlib.pyplot as plt
    plt.scatter(A[0,:],A[1,:], c='r', alpha = .5)
    plt.scatter(B[0,:],B[1,:], c='g', alpha = .5)
    plt.scatter(C[0,:],C[1,:], c='b', alpha = .5)
    plt.axis([-3, 3, -3, 3])
    plt.savefig('figure%d'%s)
    plt.clf()
    plt.show()

def runalltests(iter = 1):
  
  for i in range(iter):
    tra = np.random.randn(3)
    rot = np.random.rand() * np.pi * 2
    rotmat = np.array([[np.cos(rot), -np.sin(rot), 0],
                        [np.sin(rot), np.cos(rot),0],
                         [0 , 0, 1.]] )
    for f in datafunctions:
      A,B,Aori = f.__call__(rot = rot, t = tra)
      R,t =  combiner.probcombine(A,B)
      normR = la.norm(R-rotmat, ord = 2)
      normt = la.norm(t - tra, ord = 1)
      print R
      print t
      print 'test iter %d %s diffR = %f difft = %f' % (i, str(f), normR, normt)
    
    
